#include <mm/page_allocator.h>
#include <mm/paging.h>

#include <utils/bitmap.h>

static uint64_t region_count = 0;
static mm_memory_region *regions = NULL;

static bitmap_t *alloc_stat = NULL;
static uint64_t total_pages = 0;
static uint64_t usable_pages = 0;
static uint64_t used_pages = 0;

static bool
efi_memtype_filter(enum efi_memory_type type) {
	switch(type) {
		case EfiBootServicesCode:
		case EfiBootServicesData:
		case EfiConventionalMemory:
		case EfiACPIReclaimMemory:
			return true;
		default:
			return false;
	}
}

static bool
efi_try_fixup_mmap_hole(efi_memory_descriptor *entry, efi_memory_descriptor *prev, mm_memory_region **rentry) {
	if(entry->physical_start == 0x100000) {
		return false;
	}

	if(entry->physical_start != prev->physical_start + (prev->number_of_pages << PAGE_4KSHIFT)) {
		// __early_panic(entry->physical_start, prev->physical_start);

		++*rentry;
		++region_count;
		mm_memory_region *region = *rentry;
		++*rentry;
		++region_count;

		region->number_of_pages = (entry->physical_start - prev->physical_start) >> PAGE_4KSHIFT;
		region->number_of_pages -= prev->number_of_pages;
		region->number_of_free_pages = 0;
		region->attributes = MR_PRESENT | MR_SKIP | MR_MEMHOLE;

		return true;
	}

	return false;
}

void mm_pallocator_init(efi_handle_t image_handle, efi_system_table *system_table) {
	uint64_t map_size = 0;
	uint64_t map_key = 0;
	uint64_t desc_size = 0;
	uint32_t desc_ver = 0;

	uintptr_t last_page_address = 0;

	efi_boot_services *bs = system_table->boot_services;

	bs->get_memory_map(&map_size, NULL, &map_key, &desc_size, &desc_ver);

	{
		efi_memory_descriptor mmap[map_size / desc_size];

		bs->get_memory_map(&map_size, mmap, &map_key, &desc_size, &desc_ver);

		for(efi_memory_descriptor *entry = mmap, *prev = mmap; map_size; map_size -= desc_size, prev = entry,
															entry = (efi_memory_descriptor *)((uintptr_t)entry + desc_size)) {
			if(entry->physical_start < 0x100000) {
				continue;
			}

			if(entry->physical_start != prev->physical_start + (prev->number_of_pages << PAGE_4KSHIFT)) {
				++region_count;
			}

			bool usable = efi_memtype_filter(entry->type);

			if(usable) {
				usable_pages += entry->number_of_pages;
			}

			if(usable != efi_memtype_filter(prev->type) || (entry->type == EfiACPIReclaimMemory) != (prev->type == EfiACPIReclaimMemory)) {
				++region_count;
			}
		}

		++region_count;
	}

	size_t buf_size = sizeof(mm_memory_region) * region_count + BITMAP_PHYSIZE(usable_pages);
	void *buf = NULL;

	bs->allocate_pages(EfiAllocateAnyPages, EfiLoaderData, (buf_size >> 12) + !!(buf_size & 0xfff), (void **)&buf);
	bs->setmem(buf, sizeof(mm_memory_region) * region_count, 0);

	regions = buf;
	alloc_stat = buf + sizeof(mm_memory_region) * region_count;

	usable_pages = 0;
	region_count = 0;

	bs->get_memory_map(&map_size, NULL, &map_key, &desc_size, &desc_ver);

	{
		mm_memory_region *rentry = regions;
		efi_memory_descriptor mmap[map_size / desc_size];

		bs->get_memory_map(&map_size, mmap, &map_key, &desc_size, &desc_ver);

		rentry->attributes |= MR_PRESENT;

		if(efi_memtype_filter(mmap[0].type)) {
			rentry->attributes |= MR_USABLE;
		}

		if(mmap[0].type == EfiACPIReclaimMemory) {
			rentry->attributes |= MR_ACPI;
		}

		for(efi_memory_descriptor *entry = mmap, *prev = entry; map_size; map_size -= desc_size, prev = entry,
															entry = (efi_memory_descriptor *)((uintptr_t)entry + desc_size)) {
			if(entry->physical_start < 0x100000) {
				continue;
			}

			if(entry->attribute == 0) {
				break;
			}

			if(efi_try_fixup_mmap_hole(entry, prev, &rentry)) {
				rentry->attributes |= MR_PRESENT;

				if(efi_memtype_filter(entry->type)) {
					rentry->attributes |= MR_USABLE;
				}

				if(entry->type == EfiACPIReclaimMemory) {
					rentry->attributes |= MR_ACPI;
				}
			}

			last_page_address = entry->physical_start + ((entry->number_of_pages - 1) << PAGE_4KSHIFT);

			bool usable = efi_memtype_filter(entry->type);

			if(usable == !!(rentry->attributes & MR_USABLE) && (entry->type == EfiACPIReclaimMemory) == !!(rentry->attributes & MR_ACPI)) {
				rentry->number_of_pages += entry->number_of_pages;
				rentry->number_of_free_pages = rentry->number_of_pages;
			} else {
				if(rentry->attributes & MR_USABLE) {
					usable_pages += rentry->number_of_pages;
				} else {
					rentry->number_of_free_pages = 0;
				}

				++rentry;
				++region_count;

				if(!(rentry->attributes & MR_PRESENT)) {
					rentry->attributes |= MR_PRESENT;

					if(efi_memtype_filter(entry->type)) {
						rentry->attributes |= MR_USABLE;
					}

					if(entry->type == EfiACPIReclaimMemory) {
						rentry->attributes |= MR_ACPI;
					}
				}

				rentry->number_of_pages = entry->number_of_pages;
				rentry->number_of_free_pages = rentry->number_of_pages;
			}
		}

		if(rentry->attributes & MR_USABLE) {
			usable_pages += rentry->number_of_pages;
		} else {
			rentry->number_of_free_pages = 0;
		}

		++region_count;
	}

	total_pages = (last_page_address >> PAGE_4KSHIFT) + 1;

	alloc_stat->size = usable_pages;

	system_table->boot_services->exit_boot_services(image_handle, map_key);

	bitmap_set(alloc_stat, 0, alloc_stat->size, false);
}

uintptr_t mm_pallocator_alloc4k(uint64_t count) {
	if(!count) {
		return 0;
	}

	uint64_t allocation_base = 0;
	uint64_t physical_base = 0;
	for(uint64_t i = 0; i < region_count; ++i) {
		mm_memory_region *entry = &regions[i];
		uint64_t limit = entry->number_of_pages - count + 1;

		if(!(entry->attributes & MR_USABLE) || entry->number_of_free_pages < count) {
			goto skip;
		}

		for(uint64_t j = 0; j < limit; ++j) {
			uint64_t base = allocation_base + j;

			if(bitmap_is_all_set_as(alloc_stat, base, count, false)) {
				bitmap_set(alloc_stat, base, count, true);

				entry->number_of_free_pages -= count;
				used_pages += count;

				return (physical_base + j + 0x100) << 12;
			}
		}
	skip:
		physical_base += entry->number_of_pages;

		if(entry->attributes & MR_USABLE) {
			allocation_base += entry->number_of_pages;
		}
	}

	return 1;
}

void *mm_pallocator_valloc4k(uint64_t count, uint64_t map_flags) {
	uintptr_t p = mm_pallocator_alloc4k(count);
	if(!p) {
		return NULL;
	}

	int stat = mm_paging_map_memory(NULL, p, p, count, map_flags);
	if(stat < 0) {
		mm_pallocator_free4k(p, count);

		return NULL;
	}

	return (void *)p;
}

int mm_pallocator_free4k(uintptr_t address, uint64_t count) {
	if(!address || !count || address & 0xfff) {
		return -EINVAL;
	}

	uint64_t allocation_base = 0;
	uint64_t physical_base = 0;
	for(uint64_t i = 0; i < region_count; ++i) {
		mm_memory_region *entry = &regions[i];

		if(!(entry->attributes & MR_USABLE) || entry->number_of_pages < count) {
			goto skip;
		}

		if(((physical_base + entry->number_of_pages + 0x100) << 12) < address) {
			goto skip;
		}

		uint64_t off = (address >> 12) - physical_base - 0x100;
		uint64_t base = allocation_base + off;

		if(entry->number_of_pages < off + 1 + count) {
			return -EINVAL;
		}

		if(!bitmap_is_all_set_as(alloc_stat, base, count, true)) {
			return -EINVAL;
		}

		bitmap_set(alloc_stat, base, count, false);

		entry->number_of_free_pages += count;
		used_pages -= count;

		return 0;
	skip:
		physical_base += entry->number_of_pages;

		if(entry->attributes & MR_USABLE) {
			allocation_base += entry->number_of_pages;
		}
	}

	return -EFAULT;
}

int mm_pallocator_vfree4k(void *p, uint64_t count) {
	uintptr_t address = (uintptr_t)p;

	if(!address || !count || address & 0xfff) {
		return -EINVAL;
	}

	return mm_pallocator_free4k(mm_paging_get_paddr(NULL, address), count);
}

mm_memory_region *mm_pallocator_get_memory_regions(uint64_t *output_count) {
	if(output_count) *output_count = region_count;

	return regions;
}

uint64_t mm_pallocator_get_total_pages() {
	return total_pages;
}

uint64_t mm_pallocator_get_usable_pages() {
	return usable_pages;
}

uint64_t mm_pallocator_get_used_pages() {
	return used_pages;
}

void acpi_free_reclaimable_memory() {
	for(uint64_t i = 0; i < region_count; ++i) {
		if(regions[i].attributes & MR_ACPI) {
			regions[i].attributes &= ~MR_ACPI;
		}
	}

	uint32_t last_attribute = 0;
	for(mm_memory_region *entry = regions, *prev = entry; region_count; --region_count, ++entry) {
		if(last_attribute == entry->attributes) {
			prev->number_of_pages += entry->number_of_pages;
			prev->number_of_free_pages += entry->number_of_free_pages;

			entry->number_of_pages = 0;
			entry->number_of_free_pages = 0;
			entry->attributes = MR_PRESENT | MR_SKIP;
		} else {
			last_attribute = entry->attributes;
			prev = entry;
		}
	}
}